import torch.nn as nn
import torch


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
input_dim = 1
output_dim = 1


class LinearRegressionModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = self.linear(x)
        return out


model = LinearRegressionModel(input_dim, output_dim)
model.to(device)

epochs = 10000
learning_rate = 0.01
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
for epoch in range(epochs):
    epoch += 1
    inputs = torch.tensor([[1.0], [2.0], [3.0], [4.0]]).to(device)
    targets = torch.tensor([[2.0], [4.0], [6.0], [8.0]]).to(device)
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        print('epoch {}, loss {}'.format(epoch, loss.item()))
        print('w: ', model.linear.weight.item())
        print('b: ', model.linear.bias.item())


# 测试
print(model(torch.tensor([[5.0]]).to(device)).to(device))

torch.save(model.state_dict(), 'demo2model.pkl')